import Config as conf
import Util as util
import csv
import numpy as np
from datetime import datetime
from sklearn.utils import shuffle
from sklearn.preprocessing import StandardScaler
import pandas as pd
import scipy
from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer
from sklearn.preprocessing import LabelBinarizer
from sklearn.linear_model import Ridge

def get_data(is_shuffle=False, force_write=False):

    # all data's binary file path
    path = conf.DATA_CACHE_FILE_PATH

    if util.is_file_or_directory_exist(path) and not force_write:
        # log
        print("load train data from file")
        return util.load_from_b(path)

    else:
        # read raw data
        print("read train data from tsv...")
        df_train = pd.read_table(conf.TRAIN_DATA_PATH, engine='c')
        print("read test data from tsv...")
        df_test = pd.read_table(conf.TEST_DATA_PATH, engine='c')

        # y_train
        print("log1p with price")
        y_train = np.log1p(df_train['price'])

        # row number of train data
        nrow_train = df_train.shape[0]
        print("train data size:{}".format(nrow_train))

        # concat train & test data
        print("concat train & test data")
        df_all = pd.concat([df_train, df_test])

        # release memory
        del df_train

        """
        data processing
        """

        # name
        print("name ...")
        df_all['name'] = df_all['name'].fillna('unknown').astype('category')
        cv = CountVectorizer(ngram_range=(1, 2))
        X_name = cv.fit_transform(df_all['name'])

        # category_name
        print("category_name ...")
        df_all['category_name'] = df_all['category_name'].fillna('unknown').astype('category')
        cv = CountVectorizer()
        X_category_name = cv.fit_transform(df_all['category_name'])

        # brand_name
        print("brand_name ...")
        df_all['brand_name'] = df_all['brand_name'].fillna('unknown').astype('category')
        lb = LabelBinarizer(sparse_output=True)
        X_brand_name = lb.fit_transform(df_all['brand_name'])

        # shipping & item_condition_id
        print("shipping & item_condition_id ...")
        X_shipping_and_item_condition_id = scipy.sparse.csr_matrix(pd.get_dummies(df_all[['item_condition_id', 'shipping']], sparse=True))

        # item_description
        print("item_description ...")
        df_all['item_description'].fillna('unknown', inplace=True)
        tfidf = TfidfVectorizer(ngram_range=(1,3), max_features=50000, stop_words='english')
        X_item_description = tfidf.fit_transform(df_all['item_description'])

        # stack all features
        print("stack all features...")
        X_all = scipy.sparse.hstack([X_name, X_brand_name, X_shipping_and_item_condition_id, X_category_name, X_item_description]).tocsr()

        # save to disk
        print("save to disk ...")
        util.save_to_b(path, (nrow_train, X_all, y_train), force_write)

        return nrow_train, X_all, y_train

# get result csv file
def get_result_csv_file(df_test, file='preds.csv'):
    df_test['test_id', 'price'].to_csv(file, index=False)
